# Evaluate of SIFA
import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import numpy as np
# import matplotlib

# matplotlib.use('Agg')
# import matplotlib.pyplot as plt
from metrics import create_visual_anno
import cv2
import albumentations as A
# from tqdm import tqdm



data_root='/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA/useData4label/target'
data_list=[os.path.join(data_root,f) for f in os.listdir(data_root)]
save_path="/home/liukai/projects/SIFA/output/showPngtarget"
data_list=sorted(data_list)
def norm_01(image):
    mn = np.min(image)
    mx = np.max(image)
    image = (image - mn) / (mx - mn).astype(np.float32)
    return image

def save_img(image):
    image = norm_01(image)
    image = (image * 255).astype(np.uint8)
    return image
def load_npz(path):
    img = np.load(path)['arr_0']
    gt = np.load(path)['arr_1']
    return img, gt
transform = A.OneOf(
            [
                #A.RandomSizedCrop(p=1,min_max_height=(230,250),height=256,width=256)
                #A.Affine(p=1),#仿射变换，SIFA论文用到的
                # A.Rotate(limit=50,p=0.5)
                A.GaussNoise(p=1)
            ])
for data_name in data_list:
    img, gt = load_npz(data_name)
    print(set(gt.flatten()))

    gt_vis = create_visual_anno(gt)
    gt_vis=transform(image=gt_vis)['image']
    print(gt_vis.shape)
    if (not os.path.exists(save_path)):
        os.makedirs(save_path)
        print(save_path)
    name=os.path.basename(data_name.split('.')[0])

    cv2.imwrite('{}/gt-{}.jpg'.format(save_path, name), gt_vis)
    print(data_name)


'''
    脾（spleen）：1         红色
    右肾（right kidney）：2  绿色
    左肾（left kidney）：3   黄色
    肝：4                   浅蓝色
'''
